package org.langcom.crypt;

import java.security.AlgorithmParameterGenerator;
import java.security.AlgorithmParameters;
import java.security.InvalidAlgorithmParameterException;
import java.security.InvalidKeyException;
import java.security.KeyFactory;
import java.security.KeyPair;
import java.security.KeyPairGenerator;
import java.security.NoSuchAlgorithmException;
import java.security.PublicKey;
import java.security.spec.InvalidKeySpecException;
import java.security.spec.InvalidParameterSpecException;
import java.security.spec.X509EncodedKeySpec;

import javax.crypto.BadPaddingException;
import javax.crypto.Cipher;
import javax.crypto.IllegalBlockSizeException;
import javax.crypto.KeyAgreement;
import javax.crypto.NoSuchPaddingException;
import javax.crypto.SecretKey;
import javax.crypto.spec.DHParameterSpec;

public class DHAlice {

	private KeyAgreement aliceKeyAgree;
	private SecretKey aliceDesKey;

	public byte[] getAlicePublicKey() throws NoSuchAlgorithmException, InvalidParameterSpecException, InvalidAlgorithmParameterException, InvalidKeyException {
		String mode = "GENERATE_DH_PARAMS";
		DHParameterSpec dhSkipParamSpec;
		if (mode.equals("GENERATE_DH_PARAMS")) {
			// Some central authority creates new DH parameters
			System.out.println("Creating Diffie-Hellman parameters (takes VERY long) ...");
			AlgorithmParameterGenerator paramGen = AlgorithmParameterGenerator.getInstance("DH");
			paramGen.init(512);
			AlgorithmParameters params = paramGen.generateParameters();
			dhSkipParamSpec = (DHParameterSpec) params.getParameterSpec(DHParameterSpec.class);
		} else {
			// use some pre-generated, default DH parameters
			System.out.println("Using SKIP Diffie-Hellman parameters");
			dhSkipParamSpec = new DHParameterSpec(DHCommon.skip1024Modulus, DHCommon.skip1024Base);
		}
		/*
		 * Alice creates her own DH key pair, using the DH parameters from above
		 */
		System.out.println("ALICE: Generate DH keypair ...");
		KeyPairGenerator aliceKpairGen = KeyPairGenerator.getInstance("DH");
		aliceKpairGen.initialize(dhSkipParamSpec);
		KeyPair aliceKpair = aliceKpairGen.generateKeyPair();
		// Alice creates and initializes her DH KeyAgreement object
		System.out.println("ALICE: Initialization ...");
		aliceKeyAgree = KeyAgreement.getInstance("DH");
		aliceKeyAgree.init(aliceKpair.getPrivate());
		// Alice encodes her public key, and sends it over to Bob.
		byte[] alicePubKeyEnc = aliceKpair.getPublic().getEncoded();
		return alicePubKeyEnc;
	}

	public void aliceObtainDesKey(byte[] bobPubKeyEnc) throws NoSuchAlgorithmException, InvalidKeySpecException, InvalidKeyException, IllegalStateException {
		KeyFactory aliceKeyFac = KeyFactory.getInstance("DH");
		X509EncodedKeySpec x509KeySpec = new X509EncodedKeySpec(bobPubKeyEnc);
		PublicKey bobPubKey = aliceKeyFac.generatePublic(x509KeySpec);
		System.out.println("ALICE: Execute PHASE1 ...");
		aliceKeyAgree.doPhase(bobPubKey, true);
		// byte[] aliceSharedSecret = aliceKeyAgree.generateSecret();
		// int aliceLen = aliceSharedSecret.length;
		// byte[] bobSharedSecret = new byte[aliceLen];
		// aliceKeyAgree.doPhase(bobPubKey, true);
		aliceDesKey = aliceKeyAgree.generateSecret("DES");
	}

	public String decrypt(byte[] ciphertext) throws NoSuchAlgorithmException, NoSuchPaddingException, InvalidKeyException, IllegalBlockSizeException, BadPaddingException {
		Cipher aliceCipher = Cipher.getInstance("DES/ECB/PKCS5Padding");
		aliceCipher.init(Cipher.DECRYPT_MODE, aliceDesKey);
		byte[] recovered = aliceCipher.doFinal(ciphertext);
		return new String(recovered);
	}

	public byte[] encrypt(String text) throws NoSuchAlgorithmException, NoSuchPaddingException, InvalidKeyException, IllegalBlockSizeException, BadPaddingException {
		/*
		 * DES in ECB mode
		 */
		Cipher aliceCipher = Cipher.getInstance("DES/ECB/PKCS5Padding");
		aliceCipher.init(Cipher.ENCRYPT_MODE, aliceDesKey);
		byte[] ciphertext = aliceCipher.doFinal(text.getBytes());
		return ciphertext;
	}
}
